#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sat Nov 23 09:40:52 2019

@author: fu
"""
import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
from numpy import array, argmax, newaxis, std, dstack, savetxt
from tensorflow.keras.preprocessing import sequence
from sklearn.preprocessing import LabelEncoder, OneHotEncoder, StandardScaler, MinMaxScaler
from sklearn import metrics
from sklearn.metrics import roc_auc_score, auc, accuracy_score
from sklearn.utils.class_weight import compute_class_weight
from sklearn.model_selection import train_test_split
from sklearn.utils import class_weight, shuffle
from sklearn.decomposition import NMF
from pandas import read_csv
from tensorflow.keras.preprocessing.text import one_hot, Tokenizer
from tensorflow.keras.models import Sequential, Model, load_model
from tensorflow.keras.layers import Attention, Dense, Input, Embedding, Flatten, concatenate,\
Dropout, Activation, add, Dot, LSTM, Subtract, Lambda, RepeatVector, multiply
from tensorflow.keras.utils import to_categorical
from tensorflow.keras import backend as K
from tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint
from tensorflow.keras.constraints import non_neg
from tensorflow.keras.utils import get_custom_objects
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.initializers import glorot_uniform, RandomNormal, GlorotNormal

from tensorflow.keras.utils import model_to_dot
from matplotlib import pyplot
from tensorflow.keras.regularizers import l1_l2, l1, l2
from numpy.random import seed
from sklearn.preprocessing import OrdinalEncoder
import seaborn as sns

seed(1)
tf.random.set_seed(42)
"""
READ DATA 
"""
i2501_action = read_csv('/Users/fu/data/i2501_action.csv',header=None)
i2501_score = read_csv('/Users/fu/data/i2501_score.csv',header=None)
i2502_action = read_csv('/Users/fu/data/i2502_action.csv',header=None)
i2502_score = read_csv('/Users/fu/data/i2502_score.csv',header=None)
i0701_action = read_csv('/Users/fu/data/i0701_action.csv',header=None)
i0701_score = read_csv('/Users/fu/ddata/i0701_score.csv',header=None)
i0703_action = read_csv('/Users/fu/data/i0703_action.csv',header=None)
i0703_score = read_csv('/Users/fu/data/i0703_score.csv',header=None)
ps_composite = read_csv('/Users/fu/data/ps_composite.csv',header=None)

"""
DATA PREPROCESSING
"""
## response matrix
scores = pd.concat([i2501_score,i2502_score,i0701_score,i0703_score], axis =1)
scores.columns = ['0', '1','2','3']
scores['s_id'] = scores.reset_index().index
response = pd.melt(scores, id_vars=['s_id'], value_vars=['0', '1','2','3'])
response.columns = ['s_id', 'i_id','score']

"""
ACTION ID TO NUMBER MAPPING
"""
actions_og = pd.concat([i2501_action,i2502_action,i0701_action,i0703_action], axis =0)
n = len(actions_og.melt()['value'].value_counts())
maxlen = len(actions_og.columns)
keys = actions_og.melt()['value'].value_counts().index
act2index = dict(zip(keys, range(1, 302)))
def f(x):
    return x.map(act2index)
actions = actions_og.apply(f, axis=1) 
actions = actions.fillna(0)

"""
STUDENT ID TO NUMBER MAPPING
"""
stu2index = pd.DataFrame(columns=['stu_id','index'])
stu2index["stu_id"] = ps_composite[5]
stu2index["index"] = scores['s_id']
stu2index = stu2index.drop_duplicates().sort_values(by=['index'])  

## shuffle the data 
#response, actions, ps_composite = shuffle(response, actions,  ps_composite, random_state = 2)

"""
TRAIN/TEST SPLIT & SCALE TIME
"""
test_size = 0.2
response_train, response_test, actions_train, actions_test = train_test_split(response, actions, test_size = test_size, random_state=502)
scores_train, scores_test = train_test_split(scores, test_size = test_size, random_state=502)

np.random.seed(1) 
item_id_train = response_train.i_id.values.astype(np.int)
stu_id_train = response_train.s_id.values.astype(np.int)
item_id_test = response_test.i_id.values.astype(np.int)
stu_id_test = response_test.s_id.values.astype(np.int)
train_labels = response_train.score.values.astype(np.int)
test_labels = response_test.score.values.astype(np.int)
#response_train, actions_train = shuffle(response_train, actions_train, random_state=1)

"""
HYPERPARAMETERS
"""
epochs = 60
n_students, n_items, n_actions = len(response.s_id.unique()), len(response.i_id.unique()), n
n_latent_factors = 1
batch_size = 64
validation_split = 0.2
act_length = maxlen


"""
MODEL 
"""
def log_model(n_actions, act_length, act_output_dim = 50, layers= [3],
              nodes_lstm = 16, n_latent_factors = n_latent_factors, relu_units = 1,
              reg_w = 0, p_dropout = 0):
    
    # Input Features
    item_input = Input(shape=[1], name='Item')
    stu_input = Input(shape=[1], name='Student')
    act_input = Input(shape=(act_length,), name='Action')
    
    # Student and Item Embedding    
    item_embedding = Embedding(input_dim = n_items, output_dim = n_latent_factors, input_length = 1,
                               embeddings_regularizer = l2(0.),embeddings_initializer = glorot_uniform(seed=1), 
                               name='Item-Embedding')(item_input)
    stu_embedding  = Embedding(input_dim = n_students, output_dim = n_latent_factors, input_length = 1, #embeddings_constraint = non_neg(), 
                               embeddings_regularizer = l2(0.), embeddings_initializer = glorot_uniform(seed=1), 
                               name='Student-Embedding')(stu_input)    
    
    # CF of Items and Students
    item_vec = Flatten(name = 'FlattenItems')(item_embedding)
    stu_vec = Flatten(name = 'FlattenStudents')(stu_embedding)
    comb = concatenate([item_vec, stu_vec])
    n_layer = len(layers) 
    for i in range(n_layer):
        dropout = Dropout(p_dropout)
        dense = Dense(layers[i], kernel_initializer = glorot_uniform(seed = 1), 
                      kernel_regularizer = l2(reg_w), 
                      activation='relu', name="layer%d" %i)
        comb = dropout(comb)
        comb = dense(comb)    

    # Action Embedding
    embed = Embedding(
            input_dim = n_actions+1, 
            output_dim = act_output_dim, 
            input_length= act_length, 
            embeddings_regularizer = l1_l2(l1=0, l2=0.), #
            embeddings_initializer = glorot_uniform(seed = 1), 
            mask_zero = True,
            name = 'Action-Embedding'
    )   
    act_embedding = embed(act_input) 
    mask = embed.compute_mask(act_input) 
       
    # LSTM
    act = LSTM(nodes_lstm, kernel_initializer = glorot_uniform(seed=1), kernel_regularizer = l2(reg_w), 
               return_sequences=True, name="act")(act_embedding, mask = mask)
    act = Dropout(p_dropout)(act)

    # Attention
    attend_weight = Attention(name='Attention')(act, mask = mask)
    attend_weight_expand = Lambda(lambda x: K.expand_dims(x))(attend_weight)
    attend_hidden = multiply([act, attend_weight_expand])
    attend_hidden = Lambda(lambda x: K.sum(x, axis=1), name="Weighted-Actions")(attend_hidden)
    attend_hidden = Dense(relu_units, activation=None, kernel_initializer = glorot_uniform(seed=1), name = 'Representation')(attend_hidden)
    
    # Concatenation and Final Prediction
    prod = concatenate([attend_hidden, comb])
    prod = Dense(1, activation='sigmoid', kernel_initializer = glorot_uniform(seed=1), name='prediction')(prod)
  
    model = Model([item_input, stu_input, act_input], prod)
    model.compile(optimizer = tf.keras.optimizers.Adam(learning_rate = 0.001), loss = 'binary_crossentropy', metrics = ["BinaryAccuracy", "AUC", "mae", tf.keras.metrics.RootMeanSquaredError(name='rmse')])
    model.summary()
    return model

"""
TRAINING AND TESTING 
"""
model = log_model(n_actions, act_length)
model.fit([item_id_train, stu_id_train, actions_train], train_labels, epochs = 80, batch_size = batch_size, validation_split = validation_split, shuffle = True)#, callbacks = [es,mc])
model.evaluate([item_id_test, stu_id_test, actions_test], test_labels, batch_size = batch_size)

"""
GET ATTENTION WEIGHTS 
"""
new_model = Model(inputs= model.input, outputs = 
                  [model.get_layer(name='Action').output, 
                   model.get_layer(name='Attention').output,
                   model.get_layer(name='Representation').output,
                   model.get_layer(name='Item').output,
                   model.get_layer(name='Student').output])
actions, att_weights, representations, items, students = new_model.predict([item_id_train, stu_id_train, actions_train])

"""
AVERAGE ACTION WEIGHTS 
"""
s = []
for i in range(0, len(actions)):
    ids = actions[i]
    wgts = att_weights[i]
    for j in range(0, len(ids)):
        act = ids[j]
        wgt = wgts[j]
        s.append([act,wgt])
s = np.asarray(s)
s = pd.DataFrame(s, columns = ['act','wgt'])
act2wgts = s.groupby('act', as_index=False).mean()
index2act = {v: k for k, v in act2index.items()} 
act2wgts['act'] = act2wgts['act'].map(index2act)

"""
ACTION WEIGHTS VISUALIZATION
"""
def f(x):
    return x.map(index2act)
actions = pd.DataFrame(actions)
actions = actions.apply(f, axis=1) 

### Strategy Wrong, Outcome Incorrect
x = 242
act_name = np.squeeze(actions[x:x+1].values)
act_name = act_name[0:9]
act_imp = np.squeeze(att_weights[x:x+1])
act_imp = act_imp[act_imp != 0]
df = pd.DataFrame({"Importance": act_imp},
                  index = act_name)
plt.figure(figsize = (6,3))
sns.heatmap(df, cmap="Reds")
plt.tight_layout()
plt.savefig('heatmap_Str_Wrg_Out_Wrg.png', dpi=400)


### Strategy Correct, Outcome Incorrect
x = 35
act_name = np.squeeze(actions[x:x+1].values)
act_name = act_name[0:14]
act_imp = np.squeeze(att_weights[x:x+1])
act_imp = act_imp[act_imp != 0]
df = pd.DataFrame({"Importance": act_imp},
                  index = act_name)
plt.figure(figsize = (6,3))
sns.heatmap(df, cmap="Reds")
plt.tight_layout()
plt.savefig('heatmap_Str_Rgt_Out_Wrg.png', dpi=400)

### Strategy Correct, Outcome Correct
x = 53
act_name = np.squeeze(actions[x:x+1].values)
act_name = act_name[0:13]
act_imp = np.squeeze(att_weights[x:x+1])
act_imp = act_imp[act_imp != 0]
df = pd.DataFrame({"Importance": act_imp},
                  index = act_name)
plt.figure(figsize = (6,3))
sns.heatmap(df, cmap="Reds")
plt.tight_layout()
plt.savefig('heatmap__Str_Rgt_Out_Rgt.png', dpi=400)




"""
CORRELATION BETWEEN PROCESS DATA REPRESENTATION AND COMPOSITE SCORE
"""
stu2rep = pd.DataFrame(columns=['stu_id','rep'])
stu2rep["stu_id"] = students[:,0]
stu2rep["rep"] = representations[:,0]
stu2rep = stu2rep.groupby('stu_id', as_index=False).mean()

ps_composite['stu_id'] = stu2index['index']
comb = pd.merge(stu2rep, ps_composite, on='stu_id',  how='left')
np.corrcoef(comb["rep"].values, comb[1].values)
comb2 = pd.merge(stu2rep, ps_composite, on='stu_id',  how='right')
comb2 = comb2.sort_values(by=['stu_id'])
savetxt('/Users/fu/Desktop/20210614_Process_Data_Special_Issue/data/process_rep.csv', comb2["rep"], delimiter=',')



"""
ITEM- AND STUDENT-SKILL ASSOCIATIONS
"""
item_embeddings = model.get_layer(name='Item-Embedding').get_weights()[0]
stu_embeddings = model.get_layer(name='Student-Embedding').get_weights()[0]
item_summary = pd.DataFrame(item_embeddings).describe()
stu_summary = pd.DataFrame(stu_embeddings).describe()
savetxt('/Users/fu/Desktop/data/item_skill.csv', item_embeddings, delimiter=',')
savetxt('/Users/fu/Desktop/data/student_skill.csv', stu_embeddings, delimiter=',')

"""
CORRELATION BETWEEN EMBEDDINGS AND COMPOSITE SCORE
"""
np.corrcoef(np.squeeze(stu_embeddings, axis=1), ps_composite[1].values)










